import argparse
import datetime
import os
import pathlib
from functools import partial

import numpy as np
import torch

from args import get_args_parser
from scene.atm_opt.ATM_OPT import ATM_OPT
from scene.atm_opt.model.ga_atm import gpu_atm_costfunc
from scene.invest_opt.INVEST_OPT import INVEST_OPT
from scene.vrptw_opt.VRPTW_OPT import VRPTW_OPT
from tools.print_result.print_atm_result import print_atm_result
from tools.print_result.print_invest_result import print_invest_result
from tools.print_result.show_draw_route import draw_route
from tools.read_config import invest_config, atm_config, vrptw_config

if __name__ == '__main__':
    parser = argparse.ArgumentParser('DeiT training and evaluation script', parents=[get_args_parser()])
    args = parser.parse_args()
    
    np.random.seed(args.seed)
    torch.random.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    print(torch.__config__.parallel_info())

    if args.problem == 1:
        args.day = datetime.datetime.strptime(parser.parse_args().date, "%Y-%m-%d").isoweekday()
        folder_path = r'./DATA/atm_opt/atm_new/'

        DATA = atm_config.ATMLoadData(folder_path + args.date, val=args.val)

        if not args.val:
            atm_nlo = ATM_OPT(args, DATA=DATA)
            atm_nlo.applyALNS()
            print_atm_result(atm_nlo)
            exit()
        else:
            convert_ori_data = np.array(
                [(DATA.w_hist[1:] / 10000 / DATA.max_cap * (1 << args.bit_precision)).astype(np.uint8)]).reshape(
                1, 6, -1)
            args.val = True
            ori_fun = partial(gpu_atm_costfunc,
                              D=torch.from_numpy(DATA.D.astype(np.int32)).to(args.device),
                              aver_1=torch.from_numpy(DATA.aver_1).to(args.device),
                              max_cap=torch.from_numpy(DATA.max_cap.astype(np.int32)).to(args.device),
                              w_0=torch.from_numpy(DATA.w_0).to(args.device),
                              cutoff_list=torch.from_numpy(DATA.cutoff_list).to(args.device),
                              I_0=torch.from_numpy(DATA.I_0.astype(np.int32)).to(args.device),
                              avg_D2=torch.from_numpy(DATA.avg_D2).to(args.device),
                              C=0,
                              std_D2=torch.from_numpy(DATA.std_D2).to(args.device),
                              cutoff_limit=torch.from_numpy(DATA.cutoff_limit).to(args.device),
                              deposit_history=torch.from_numpy(DATA.h_deposit_history[-7:]).to(args.device),
                              withdraw_history=torch.from_numpy(DATA.h_withdraw_history[-7:]).to(args.device),
                              args=args
                              )
            real = ori_fun(
                torch.from_numpy(convert_ori_data).to(
                    args.device))
            args.val = False
            print(real[-1])
            val_day = args.day
            w_ = convert_ori_data.copy()
            I, cutoff_list = None, None
            for i in range(6):
                args.day = (val_day + i - 1) % 7 + 1
                cur_d = pathlib.Path(os.path.join(
                    folder_path,
                    (datetime.datetime.strptime(os.path.basename(args.date), "%Y-%m-%d") + datetime.timedelta(
                        i)).strftime('%Y-%m-%d'),
                )).as_posix()
                DATA = atm_config.ATMLoadData(cur_d)
                if i > 0:
                    DATA.w_0 = w_0.T
                    DATA.I_0 = I.T
                    DATA.cutoff_list = cutoff_list.T

                atm_nlo = ATM_OPT(args, DATA=DATA)
                atm_nlo.applyALNS()
                print_atm_result(atm_nlo)

                w_[:, i] = atm_nlo.w[:, 0].cpu().numpy()
                args.val = True
                args.day = val_day
                val = ori_fun(torch.from_numpy(w_).to(args.device))
                args.val = False

                w_0, I, cutoff_list = val[0][:, i + 1].cpu().numpy(), val[1][:, i].cpu().numpy(), val[2][:,
                                                                                                  i].cpu().numpy(),

            print(real[-1])
            print(torch.sum(real[1]) * 7.5e-5)
            print(val[-1])
            print(torch.sum(val[1]) * 7.5e-5)




    elif args.problem == 2:
        folder_path = './DATA/vrptw_opt'
        data = vrptw_config.VRPTWLoadData(args, folder_path)

        vrptw_opt = VRPTW_OPT(args, data=data)
        curcost, _ = vrptw_opt.compute()

        # 表示方案的车辆路径，不同车辆的路径用0分割
        plan = vrptw_opt.a
        # 表示为每个网点的经纬度
        location = data.location

        print("输出图片结果")
        draw_route(plan, location)
        print("线路成本")
        print(vrptw_opt.curcost)

    elif args.problem == 3:
        # 读取股票文件代码配置文件，stock.config存储的是股票的代码
        config_path = './DATA/invest_opt/stock.config'
        cache_path = './DATA/invest_opt/invest.cache'
        # 设置对应股票训练数据的时间范围
        start_date = '20200101'
        end_date = '20220101'

        opening_price = invest_config.getStockData(start_date=start_date,
                                                   end_date=end_date, config_path=config_path, cache_path=cache_path)

        # 传入模型训练参数，初始化模型
        investment_nlo = INVEST_OPT(args, opening_price)

        # 实现模型的训练与预测
        investment_nlo.predict()
        plan, value = investment_nlo.getRec()

        print("线路成本")
        print(value)
        print("输出规划结果")
        print_invest_result(plan / np.sum(plan))

    else:
        print('not vaild')
